package com.idanchuang.component.businesslog.es.repository;

import com.alibaba.fastjson.JSON;
import com.idanchuang.component.businesslog.es.annotations.Document;
import com.idanchuang.component.businesslog.es.config.BusinessLogConfig;
import com.idanchuang.component.businesslog.es.document.AbstractDocument;
import com.idanchuang.component.businesslog.es.util.DocumentUtil;
import com.idanchuang.component.businesslog.es.util.PageResponse;
import com.zoi7.component.core.base.BaseClass;
import com.zoi7.component.core.exception.ArgumentsInvalidException;
import com.zoi7.component.core.exception.NotFoundException;
import org.elasticsearch.action.DocWriteResponse;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.ClearScrollRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchScrollRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;
import org.elasticsearch.client.indices.GetIndexRequest;
import org.elasticsearch.client.indices.GetIndexResponse;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.elasticsearch.search.sort.SortBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.lang.reflect.ParameterizedType;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * @author yjy
 * @date 2019/10/24 13:16
 **/
public abstract class AbstractRepository<T extends AbstractDocument> extends BaseClass {

    private static final Logger log = LoggerFactory.getLogger(AbstractRepository.class);

    @Autowired
    protected RestHighLevelClient client;
    @Autowired
    protected BusinessLogConfig businessLogConfig;
    protected Class<T> persistentClass;

    public AbstractRepository(RestHighLevelClient client) {
        this(client, null);
    }

    public AbstractRepository() {
        this(null, null);
    }

    private AbstractRepository(RestHighLevelClient client, BusinessLogConfig businessLogConfig) {
        if (client != null) {
            this.client = client;
        }
        if (businessLogConfig != null) {
            this.businessLogConfig = businessLogConfig;
        }
        ParameterizedType parameterizedType = (ParameterizedType) getClass().getGenericSuperclass();
        this.persistentClass = (Class<T>) parameterizedType.getActualTypeArguments()[0];
    }

    /**
     * 索引文档
     *
     * @param entity 文档对象
     * @return response
     */
    public IndexResponse index(T entity) {
        String indexName = getIndexName();
        IndexRequest request = new IndexRequest();
        request.index(indexName);
        request.source(DocumentUtil.toField(entity));
        try {
            return client.index(request, RequestOptions.DEFAULT);
        } catch (IOException e) {
            throw new RuntimeException("Error while index: " + request.toString(), e);
        }
    }

    /**
     * 批量索引文档
     *
     * @param entities 文档对象集合
     * @return org.elasticsearch.action.bulk.BulkResponse
     * @throws
     * @Author sxp
     * @Date 2020/3/10 13:26
     */
    public BulkResponse batchIndex(List<T> entities) {
        String indexName = getIndexName();
        if (StringUtils.isEmpty(indexName)) {
            throw new RuntimeException(" indexName is empty");
        }

        if (CollectionUtils.isEmpty(entities)) {
            throw new RuntimeException(" data is empty");
        }

        BulkRequest request = new BulkRequest();
        entities.forEach(entity -> {
            IndexRequest indexRequest = new IndexRequest();
            indexRequest.index(indexName);
            indexRequest.source(DocumentUtil.toField(entity));
            request.add(indexRequest);
        });

        BulkResponse bulkResponse;
        try {
            bulkResponse = client.bulk(request, RequestOptions.DEFAULT);
            if (bulkResponse.hasFailures()) {
                throw new RuntimeException(bulkResponse.buildFailureMessage());
            }
        } catch (IOException e) {
            throw new RuntimeException("Error while index: " + request.toString(), e);
        }
        return bulkResponse;
    }

    /**
     * 根据文档id删除文档
     *
     * @param id 文档id
     */
    public boolean deleteById(String id) throws IOException {
        if (isBlank(id)) {
            throw new ArgumentsInvalidException("id is empty");
        }
        DeleteResponse deleteResponse;
        List<String> indexList = getRealIndices();
        if (indexList == null || indexList.isEmpty()) {
            throw new NotFoundException("index not found");
        }
        for (String indexName : indexList) {
            DeleteRequest request = new DeleteRequest(indexName, id);
            deleteResponse = client.delete(
                    request, RequestOptions.DEFAULT);
            if (deleteResponse.getResult() == DocWriteResponse.Result.NOT_FOUND) {
                continue;
            }
            if (deleteResponse.getResult() == DocWriteResponse.Result.DELETED) {
                return true;
            }
        }
        return false;
    }

    /**
     * 根据条件批量删除
     *
     * @param queryBuilder 条件
     */
    public BulkByScrollResponse deleteByQuery(QueryBuilder queryBuilder) {
        String indexName = getIndexName();
        BulkByScrollResponse bulkResponse = null;
        try {
            DeleteByQueryRequest request = new DeleteByQueryRequest(indexName);
            request.setQuery(queryBuilder);
            bulkResponse =
                    client.deleteByQuery(request, RequestOptions.DEFAULT);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        return bulkResponse;
    }

    /**
     * 根据id获取实体
     *
     * @param id id
     * @return 文档
     */
    public T getDoc(String id) {
        if (isBlank(id)) {
            return null;
        }
        try {
            List<String> indexList = getRealIndices();
            if (indexList == null || indexList.isEmpty()) {
                return null;
            }
            for (String indexName : indexList) {
                GetRequest request = new GetRequest(indexName, id);
                GetResponse getResponse = client.get(request, RequestOptions.DEFAULT);
                if (getResponse.isExists()) {
                    return map2Entity(getResponse.getSourceAsString(), id);
                } else {
                    log.warn("getResponse not exist");
                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return null;
    }

    /**
     * 获取记录 (最多只会返回10条结果, 如果需要全量数据, 请使用 getPage 或 searchAll 接口)
     *
     * @param searchSourceBuilder 查询条件
     * @return 记录
     */
    public List<T> search(SearchSourceBuilder searchSourceBuilder) {
        SearchRequest searchRequest = new SearchRequest(getIndexName());
        searchRequest.source(searchSourceBuilder);
        List<T> list = new ArrayList<>();
        try {
            SearchResponse response = client.search(searchRequest, RequestOptions.DEFAULT);
            for (SearchHit hit : response.getHits()) {
                list.add(map2Entity(hit.getSourceAsString(), hit.getId()));
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return list;
    }

    /**
     * 获取所有记录
     *
     * @param searchSourceBuilder 查询条件
     * @return 所有记录
     */
    public List<T> searchAll(SearchSourceBuilder searchSourceBuilder) {
        List<T> list = new ArrayList<>();
        String indexName = getIndexName();
        searchSourceBuilder.size(businessLogConfig.getScrollSize());
        SearchRequest searchRequest = new SearchRequest(indexName);
        searchRequest.source(searchSourceBuilder);
        searchRequest.scroll(TimeValue.timeValueSeconds(businessLogConfig.getScrollTimeout()));
        String scrollId = null;
        try {
            SearchResponse response =  this.client.search(searchRequest, RequestOptions.DEFAULT);
            if (response.getHits() != null && response.getHits().getHits().length > 0) {
                for (SearchHit hit : response.getHits()) {
                    list.add(map2Entity(hit.getSourceAsString(), hit.getId()));
                }
                while (response.getScrollId() != null && response.getHits().getHits().length > 0) {
                    scrollId = response.getScrollId();
                    SearchScrollRequest searchScrollRequest = new SearchScrollRequest(response.getScrollId());
                    searchScrollRequest.scroll(TimeValue.timeValueSeconds(businessLogConfig.getScrollTimeout()));
                    response = client.scroll(searchScrollRequest, RequestOptions.DEFAULT);
                    if (response.getHits() != null && response.getHits().getHits().length > 0) {
                        for (SearchHit hit : response.getHits()) {
                            list.add(map2Entity(hit.getSourceAsString(), hit.getId()));
                        }
                    }
                }
            }
            // 清除滚动
            if (scrollId != null) {
                //创建一个新的ClearScrollRequest
                ClearScrollRequest request = new ClearScrollRequest();
                //将滚动id添加到要清除的滚动标识符列表中
                request.addScrollId(scrollId);
                client.clearScroll(request, RequestOptions.DEFAULT);
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return list;
    }

    /**
     * 获取总记录数
     *
     * @param searchSourceBuilder 搜索条件
     * @return 总记录数
     */
    public int count(SearchSourceBuilder searchSourceBuilder) {
        CountRequest countRequest = new CountRequest(getIndexName());
        countRequest.source(searchSourceBuilder);
        try {
            CountResponse response = client.count(countRequest, RequestOptions.DEFAULT);
            return (int) response.getCount();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 获取分页数据
     *
     * @param searchSourceBuilder 搜索条件
     * @param size                每页条数
     * @return 第一页数据
     * @deprecated 此方法不支持 sort / fetchSource 等功能, 请结合 count + getPageList 来替代
     * @throws IOException e
     */
    @Deprecated
    public PageResponse<T> getPage(SearchSourceBuilder searchSourceBuilder, int pageNo, int size) throws IOException {
        pageNo = Math.max(pageNo, 1);
        size = Math.max(size, 1);
        int from = (pageNo - 1) * size;
        // 获取总数
        int count = count(searchSourceBuilder);
        if (count <= from) {
            return new PageResponse<T>(count, pageNo, size);
        }
        List<T> list = getPageList(searchSourceBuilder, pageNo, size);
        return new PageResponse<>(list, count, pageNo, size);
    }

    /**
     * 获取分页数据
     *
     * @param searchSourceBuilder 搜索条件
     * @param size                每页条数
     * @return 第一页数据
     * @throws IOException e
     */
    public List<T> getPageList(SearchSourceBuilder searchSourceBuilder, int pageNo, int size) throws IOException {
        pageNo = Math.max(pageNo, 1);
        size = Math.max(size, 1);
        int from = (pageNo - 1) * size;
        String indexName = getIndexName();
        List<T> list;
        // 如果 from + size 没有超过 max_result_window 的限制, 则直接走 from/size 分页查询
        if (from + size < businessLogConfig.getMaxResultWindow()) {
            list = getListNormal(indexName, searchSourceBuilder, from, size, null);
        }
        // 使用 searchAfter 进行深度分页查询
        else {
            // 如果未指定排序, 则使用 _id 字段进行排序
            List<SortBuilder<?>> sorts = searchSourceBuilder.sorts();
            if (sorts == null || sorts.stream().noneMatch((builder) -> "_id".equals(((FieldSortBuilder)sorts.get(0)).getFieldName()))) {
                searchSourceBuilder.sort("_id", SortOrder.DESC);
            }
            int currentSize = 0;
            int scrollSize = businessLogConfig.getSearchAfterSize();
            FetchSourceContext fetchSourceContext = searchSourceBuilder.fetchSource();
            Object[] currentSort = null;
            while (currentSize < from) {
                int thisScroll = Math.min(scrollSize, from - currentSize);
                searchSourceBuilder.size(thisScroll);
                if (currentSort != null) {
                    searchSourceBuilder.searchAfter(currentSort);
                }
                searchSourceBuilder.fetchSource(new String[]{"id"}, null);
                SearchRequest searchRequest = new SearchRequest(indexName);
                searchRequest.source(searchSourceBuilder);
                SearchResponse response = this.client.search(searchRequest, RequestOptions.DEFAULT);
                SearchHit[] hits = response.getHits().getHits();
                if (hits == null || hits.length < thisScroll) {
                    throw new ArgumentsInvalidException("No such records");
                }
                currentSort = hits[hits.length - 1].getSortValues();
                currentSize += thisScroll;
            }
            searchSourceBuilder.fetchSource(fetchSourceContext);
            list = getListNormal(indexName, searchSourceBuilder, 0, size, currentSort);
        }
        return list;
    }

    /**
     * 查询具体数据
     * @param indexName 索引
     * @param searchSourceBuilder 查询条件
     * @param from offset
     * @param size 条数
     * @param searchAfter searchAfter
     * @return 列表
     * @throws IOException e
     */
    private List<T> getListNormal(String indexName, SearchSourceBuilder searchSourceBuilder, int from, int size, Object[] searchAfter) throws IOException {
        List<T> list = new ArrayList<>();
        searchSourceBuilder.from(from);
        searchSourceBuilder.size(size);
        if (searchAfter != null && searchAfter.length > 0) {
            searchSourceBuilder.searchAfter(searchAfter);
        }
        SearchRequest searchRequest = new SearchRequest(indexName);
        searchRequest.source(searchSourceBuilder);
        SearchResponse response = this.client.search(searchRequest, RequestOptions.DEFAULT);
        if (response.getHits() != null) {
            for (SearchHit hit : response.getHits()) {
                list.add(map2Entity(hit.getSourceAsString(), hit.getId()));
            }
        }
        return list;
    }

    /**
     * 获取分页数据
     *
     * @param searchSourceBuilder 搜索条件
     * @param pageNo              页号
     * @param size                每页条数
     * @param allCount            总记录数 (通过count接口获取)
     * @return 第一页数据
     * @throws IOException e
     */
    public PageResponse<T> getPage(SearchSourceBuilder searchSourceBuilder, int pageNo, int size, int allCount) throws IOException {
        List<T> list = new ArrayList<>();
        if (allCount > 0) {
            list = getPageList(searchSourceBuilder, pageNo, size);
        }
        return new PageResponse<>(list, allCount, pageNo, size);
    }

    /**
     * 获取所有真实索引列表
     * @return 真实索引列表
     * @throws IOException e
     */
    protected List<String> getRealIndices() throws IOException {
        String index = getIndexName();
        GetIndexRequest request = new GetIndexRequest(index + "-*");
        GetIndexResponse response = this.client.indices().get(request, RequestOptions.DEFAULT);
        return new ArrayList<>(Arrays.asList(response.getIndices()));
    }

    /**
     * 数据转换为实体
     *
     * @param sourceString 查询结果
     * @param id 文档id
     * @return 实体
     */
    protected T map2Entity(String sourceString, String id) {
        T t = JSON.parseObject(sourceString, this.persistentClass);
        t.setId(id);
        return t;
    }

    /**
     * 从实体上获取索引信息
     *
     * @return ..
     */
    protected String getIndexName() {
        Document document = this.persistentClass.getAnnotation(Document.class);
        if (document == null) {
            throw new RuntimeException("Document must be annotation");
        }
        return document.indexName();
    }
}
